import torch
from torch.utils.data import Dataset

from pathlib import Path
import numpy as np

from PIL import Image
from .VideoAugmentation import random_crop, random_flip


class Vimeo90K(Dataset):
    """Load a video folder database. Training and testing video clips
    are stored in a directorie containing mnay sub-directorie like Vimeo90K Dataset:

    .. code-block::

        - rootdir/
            train.list
            test.list
            - sequences/
                - 00010/
                    ...
                    -0932/
                    -0933/
                    ...
                - 00011/
                    ...
                - 00012/
                    ...

    training and testing (valid) clips are withdrew from sub-directory navigated by
    corresponding input files listing relevant folders.

    This class returns a set of three video frames in shape T, C, H, W

    Args:
        root (string): root directory of the dataset
        frame_num(int): frame count
        rnd_frame_group: random take [frames] frames from extracted sequences (random start index)
        interval: take [frames] frames at [interval], default: 1 for take consecutive frames
        split (string): split mode ('train' or 'test')
    """

    def __init__(
        self,
        root, split_file,
        frame_num = 2, rnd_frame_group = True,
        interval = 1,
        im_width = 256, im_height = 256,
    ):
    
        splitfile = Path(f"{root}/{split_file}")
        splitdir = Path(f"{root}/sequences")

        if not splitfile.is_file():
            raise RuntimeError(f'Missing file "{splitfile}"')

        if not splitdir.is_dir():
            raise RuntimeError(f'Missing directory "{splitdir}"')

        with open(splitfile, "r") as f_in:
            self.sample_folders = [Path(f"{splitdir}/{f.strip()}") for f in f_in]

        self.interval = interval
        self.rnd_frame_group = rnd_frame_group
        self.frame_num = frame_num

        self.im_width, self.im_height = im_width, im_height

    def __getitem__(self, index):

        sample_folder = self.sample_folders[index]
        samples = sorted(f for f in sample_folder.iterdir() if f.is_file())

        max_index = len(samples) - ((self.frame_num - 1) * self.interval)
        start_index = np.random.randint(0, max_index) if self.rnd_frame_group else 0

        # take frames at [interval] starting at Frame [start_index]
        frame_paths = [
            samples[start_index + i*self.interval] 
            for i in range(self.frame_num)
        ]
        
        frames = np.concatenate(
            [np.asarray(Image.open(p).convert("RGB")).transpose(2, 0, 1) for p in frame_paths], 
            axis = 0
        )
        # T * C, H, W
        frames = torch.Tensor(frames / 255.0)

        if self.rnd_frame_group:
            frames = random_crop(frames, (self.im_height, self.im_width))
            frames = random_flip(frames)
            
        else:
            frames = frames[:, :self.im_height, :self.im_width]

        # T, C, H, W
        frames = frames.reshape((self.frame_num, 3, self.im_height, self.im_width))

        return frames

    def __len__(self):
        return len(self.sample_folders)





